package jtg.generator;

import jtg.generator.bean.ConverJimpleVar;
import jtg.generator.bean.InvokeStmt;
import jtg.graphics.SootCFG;
import jtg.solver.ExpressionException;
import jtg.solver.Z3Solver;
import jtg.visualizer.Visualizer;
import soot.Body;
import soot.Local;
import soot.Unit;
import soot.jimple.internal.JAssignStmt;
import soot.jimple.internal.JIfStmt;
import soot.jimple.internal.JReturnStmt;
import soot.toolkits.graph.UnitGraph;

import java.io.File;
import java.util.*;

public class SimpleGenerator {

    private String clsPath;
    private String clsName;
    private String mtdName;
    private UnitGraph ug;
    private Body body;

    private boolean longTag = false;
    private boolean floOrDouTag = false;
    private boolean stringTag = false;
    private boolean boolTag = false;

    // 出现强制类型转换的变量
    private ArrayList<ConverJimpleVar> converJimpleVars = new ArrayList<>();
    // 出现方法调用的Jimple语句， 和去掉调用方法后的路径约束语句
    private ArrayList<InvokeStmt> invokeStmts = new ArrayList<>();


    public SimpleGenerator(String className, String methodName) {
        String defaultClsPath = System.getProperty("user.dir") + File.separator + "target" + File.separator + "classes";
        new SimpleGenerator(defaultClsPath, className, methodName);
    }

    public SimpleGenerator(String classPath, String className, String methodName) {
        clsPath = classPath;
        clsName = className;
        mtdName = methodName;
        // 获取对应方法的活动方法体（Soot里的类）
        body = SootCFG.getMethodBody(clsPath, clsName, mtdName);
        // 获取完整的控制流图。UnitGraph类
        ug = SootCFG.getMethodCFG(clsPath, clsName, mtdName);
        knowParasType();
    }

    private void knowParasType() {
        for (Local para : body.getParameterLocals()) {
            if (para.getType().toString().equals("long"))
                longTag = true;
            else if (para.getType().toString().equals("float") || para.getType().toString().equals("double")) {
                floOrDouTag = true;
            } else if (para.getType().toString().equals("boolean")) {
                boolTag = true;
                ConverJimpleVar converJimpleVar = new ConverJimpleVar(para.getName().toString(), "boolean", "int");
                converJimpleVars.add(converJimpleVar);
            } else if (para.getType().toString().equals("java.lang.String")) {
                stringTag = true;
            }
        }
    }

    public void drawCFG(String graphName, boolean indexLabel) {
        Visualizer.printCFGDot(graphName, ug, indexLabel);
    }


    // 识别并返回给定方法体中的 Jimple 特有的局部变量。
    // 通常用于表示在字节码级别捕获的中间值或操作数。
    // 不是特有的局部变量：与Java源代码中的变量一一对应。可以是方法参数、类成员变量或临时变量等
    // 在某些情况下，分析 Jimple 特有的局部变量可能对于理解程序的结构和执行过程是有帮助的。
    private List<Local> getParJVars() {
        //Jimple自身增加的Locals，不是被测代码真正的变量
        ArrayList<Local> jimpleVars = new ArrayList<Local>();
        for (Local l : body.getLocals()) {
            // 对于每个局部变量，检查其字符串表示形式是否以 $ 开头。
            // 如果是，说明这是 Jimple 特有的局部变量，将其添加到 jimpleVars 列表中。
            if (l.toString().startsWith("temp$")) {
//            if (l.toString().contains("$")) {
                jimpleVars.add(l);
            }
        }
        return jimpleVars;
    }

    // 获取方法的参数
    private List<Local> getParameter() {
        ArrayList<Local> paras = new ArrayList<Local>();
        for (Local para : body.getParameterLocals()) {
            paras.add(para);
//            System.out.print(para.toString());
        }
//        System.out.println();
        return paras;
    }

    public List<String> generate() {

        // 存储路径
        List<Unit> path = null;
        // 存储生成的测试用例
        ArrayList<String> testSet = null;
        // 存储路径约束
        String pathConstraint = "";

        System.out.println("============================================================================");
        System.out.println("Generating test case inputs for method: " + clsName + "." + mtdName + "()");
        System.out.println("============================================================================");
        try {
            testSet = new ArrayList<String>();
            // 遍历所有基本块路径的头和尾
            for (Unit h : ug.getHeads())
                for (Unit t : ug.getTails()) {
                    // 获取头到尾之间的扩展基本块路径。
                    // 如果两个路径具有相同的头和尾，并且通过 ug.getExtendedBasicBlockPathBetween(h, t) 方法获取这两个路径，
                    // 该方法能确保唯一性。只能处理顺序和分支结构，循环不行
                    path = ug.getExtendedBasicBlockPathBetween(h, t);
                    System.out.println("The path is: " + path.toString());
                    // 计算路径约束，实际上就是要走该路径的条件
                    pathConstraint = calPathConstraint(path);
                    //如果路径约束为空字符串，表示路径约束为恒真，直接随机生成
                    if (pathConstraint.isEmpty()) testSet.add(randomTC(body.getParameterLocals()));
                    System.out.println("The corresponding path constraint is: " + pathConstraint);
                    // 非空，用对路径约束的解析器添加测试用例
                    // 每条路径对应一个测试输入
                    if (!pathConstraint.isEmpty())
                        testSet.add(solve(pathConstraint));
                }
        } catch (Exception e) {
            System.err.println("Error in generating test cases: ");
            System.err.println(e.toString());
        }
        // 生成的测试用例非空，输出一下
        if (!testSet.isEmpty()) {
            System.out.println("");
            System.out.println("The generated test case inputs:");
            int count = 1;
            for (String tc : testSet) {
                System.out.println("( " + count + " ) " + tc.toString());
                count++;
            }
        }
        return testSet;
    }


    private List<List<Unit>> getAllPath() {
        return null;
    }

    /**
     * 计算给定路径的路径约束
     *
     * @param path 包含单元（语句）的路径
     * @return 路径约束的字符串表示
     */
    public String calPathConstraint(List<Unit> path) {

        // 初始化路径约束和预期结果
        String pathConstraint = "";
        String expectedResult = "";

        // 用于存储赋值语句的哈希映射
        HashMap<String, String> assignList = new HashMap<>();
        // 存储带有Jimple变量的条件
        ArrayList<String> stepConditionsWithJimpleVars = new ArrayList<String>();
        // 存储条件
        ArrayList<String> stepConditions;

        // 遍历路径中的每个单元（语句）
        for (Unit stmt : path) {

            // 赋值语句
            if (stmt instanceof JAssignStmt) {
                // 先处理一下有类型转换的局部变量，将语句中的类型转换去掉，然后把他存起来
                // 等号左边和等号右边
                assignList.put(((JAssignStmt) stmt).getLeftOp().toString(), dealForTypeCon(stmt));
                continue;
            }
            // 条件语句
            if (stmt instanceof JIfStmt) {

                // 条件
                String ifstms = ((JIfStmt) stmt).getCondition().toString();
                // 下一条语句的索引
                int nextUnitIndex = path.indexOf(stmt) + 1;
                Unit nextUnit = path.get(nextUnitIndex);

                //如果ifstmt的后继语句（getTarget）不是ifstmt中goto语句，说明ifstmt中的条件为假
                if (!((JIfStmt) stmt).getTarget().equals(nextUnit))
                    ifstms = "!( " + ifstms + " )";
                else
                    ifstms = "( " + ifstms + " )";
                // WithJimpleVars并非是特有变量
                stepConditionsWithJimpleVars.add(ifstms);
                continue;
            }
            // 返回语句
            if (stmt instanceof JReturnStmt) {
                expectedResult = stmt.toString().replace("return", "").trim();
            }
        }
        System.out.println("The step conditions with JimpleVars are: " + stepConditionsWithJimpleVars);

        // 获取已经除去Jimple特有变量的路径约束
        stepConditions = replaceJvars(assignList, stepConditionsWithJimpleVars);

        // 一路上没有条件，路径约束就是空的
        if (stepConditions.isEmpty())
            return "";
        ///////////////////如果有long类型，处理long类型//////////
        if (longTag == true) dealLLong(stepConditions);
        ///////////////////如果有float类型，处理float类型//////////
        if (floOrDouTag == true) dealFloOrDou(stepConditions);
        ///////////////////如果有String类型，处理String类型//////////
        if (stringTag == true) dealString(stepConditions);
        /////////////////////////////////////////////////////
        // 获取第一个条件
        pathConstraint = stepConditions.get(0);
        int i = 1;
        while (i < stepConditions.size()) {
            // 把所有条件用‘&&’连起来
            pathConstraint = pathConstraint + " && " + stepConditions.get(i);
            i++;
        }
        // 返回约束
        //System.out.println("The path expression is: " + pathConstraint);
        return pathConstraint;
    }

    // 处理String类型
    private void dealString(ArrayList<String> stepConditions) {
        for (int i = 0; i < stepConditions.size(); i++) {
            String curCond = stepConditions.get(i);
            StringBuffer aftCond = new StringBuffer();
            // 获取当前条件语句的每个节点，是由空格分开的
            String[] nodes = curCond.split(" ");
            for(InvokeStmt invoke : invokeStmts) {
                // 是这个变量调用的方法
                if (invoke.getAfterRemInvoke().equals(curCond)) {
                    // 前面"("和变量和大于/小于/等于号
                    aftCond.append(nodes[0] + " ");
                    aftCond.append(nodes[1] + " ");
                    aftCond.append(nodes[2] + " ");
                    // 如果是compareTo函数，把对应的String参数也转换成
                    if (invoke.getInvokeMethod().contains("compareTo")) {
//                        aftCond.append(nodes[0] + " ");
//                        // 把这个条件改为多个条件(用&&相连)：String改为char[]，一个一个的处理，最后合起来
//                        // 变量的char[]类型
//                        char[] compChars = invoke.getParas().replace("\"", "").toCharArray();
//                        aftCond.append(nodes[1]+"0 " + nodes[2]+" " + compChars[0]+" ");
//                        for (int j = 1; j < compChars.length; j++) {
//                            aftCond.append("&& " + nodes[1]+j+" " + nodes[2]+" " + compChars[j]+" ");
//                        }
                        aftCond.append(invoke.getParas() + " ");

                    }
                    // 最后")"
                    aftCond.append(nodes[4]);
                }
            }
            stepConditions.remove(i);
            stepConditions.add(aftCond.toString());
        }
    }

    // 处理long类型
    private void dealLLong(ArrayList<String> stepConditions) {
        for (int i = 0; i < stepConditions.size(); i++) {
            String curCond = stepConditions.get(i);
            StringBuffer aftCond = new StringBuffer();
            if (curCond.contains("cmp")) {
                String[] nodes = curCond.split(" ");
                // 如果是long的比较
                if (nodes[2].equals("cmp")) {
                    aftCond.append(nodes[0] + " ");
                    aftCond.append(nodes[1] + " ");
                    aftCond.append(nodes[4] + " ");
                    aftCond.append(nodes[3].replace("L", "") + " ");
                    aftCond.append(nodes[6]);
                } else throw new ExpressionException("long类型处理cmp问题");
            } else throw new ExpressionException("long类型处理，有其他函数");
            stepConditions.remove(i);
            stepConditions.add(aftCond.toString());
        }
    }

    // 处理double或float类型
    private void dealFloOrDou(ArrayList<String> stepConditions) {
        for (int i = 0; i < stepConditions.size(); i++) {
            String curCond = stepConditions.get(i);
            StringBuffer aftCond = new StringBuffer();
            if (curCond.contains("cmpl")) {
                String[] nodes = curCond.split(" ");
                // 如果是long的比较
                if (nodes[2].equals("cmpl")) {
                    aftCond.append(nodes[0] + " ");
                    aftCond.append(nodes[1] + " ");
                    aftCond.append(nodes[4] + " ");
                    aftCond.append(nodes[3]+" ");
                    aftCond.append(nodes[6]);
                } else throw new ExpressionException("float类型处理cmpl问题");
            } else throw new ExpressionException("float类型处理，有其他函数");
            stepConditions.remove(i);
            stepConditions.add(aftCond.toString());
        }
    }

    // 处理强制类型转换
    private String dealForTypeCon(Unit stmt) {
        String right = ((JAssignStmt) stmt).getRightOp().toString();
        // 有强制类型转换
        if (right.contains("(") && right.contains(")")) {
            int lParenthesis = right.indexOf("(");
            int rParenthesis = matchBracket(right, lParenthesis);
            // 如果是虚拟调用方法里的参数类型，不处理
            if (right.charAt(rParenthesis+1) == '>') return right;
            ConverJimpleVar converJimpleVar = new ConverJimpleVar();
            // 截取()里的转换成的类型
            String converType = right.substring(lParenthesis+1, rParenthesis);
            // 如果截取到的内容为空字符串或含有"，那么说明这个括号不是强制类型转换，而是调用的方法或者String参数
            if (converType == null || converType.equals("") || converType.contains("\"")) return right;
            if (!converType.contains("int") && !converType.contains("long") && !converType.contains("float") &&
            !converType.contains("double") && !converType.contains("java.lang.Byte") && !converType.contains("short") &&
            !converType.contains("boolean") && !converType.contains("java.lang.String")) return right;
            converJimpleVar.setSwichType(converType);
            // 去掉这个显式的强制转换
            right = right.replace(right.substring(lParenthesis, rParenthesis+1), "").trim();
            for (Local para : body.getParameterLocals()) {
                if (right.equals(para.getName())) {
                    converJimpleVar.setVar(right);
                    converJimpleVar.setOriginalType(para.getType().toString());
                }
            }
            converJimpleVars.add(converJimpleVar);
        }
        return right;
    }

    // 处理所有语句中Jimple自定义的临时变量
    private ArrayList<String> replaceJvars(HashMap<String, String> assignList, ArrayList<String> stepConditionsWithJimpleVars) {
        // 结果：路径约束条件
        ArrayList<String> stepConditions = new ArrayList<>();
        // 获取特有的Jimple变量
        List<Local> jVars = getParJVars();

        //bug 没有考虑jVars（特有变量）为空的情况
        if (jVars.size() != 0) {
            //先把赋值语句处理一遍，因为有些赋值语句里会有jVars（特有变量）
            Iterator<Map.Entry<String, String>> iterator = assignList.entrySet().iterator();
            while (iterator.hasNext()) {
                Map.Entry<String, String> entry = iterator.next();
                //替换赋值语句里特有的Jimple变量，如果有的话
                //这里可能有一个小bug，如果存在连续的调用怎么办，可能要用一个循环
                if (entry.getValue().contains("$")) {
                    String value = entry.getValue();
                    for (Local lv : jVars) {
                        if (value.contains(lv.toString()))
                            entry.setValue(value.replace(lv.toString(), assignList.get(lv.toString()).trim()));
                    }
                }
            }

            // 对条件语句的处理
            for (int i = 0; i < stepConditionsWithJimpleVars.size(); i++) {
                String cond = stepConditionsWithJimpleVars.get(i);
                //替换条件里特有的Jimple变量
                for (Local lv : jVars) {
                    // 如果包含特有变量
                    if (cond.contains(lv.toString())) {
                        // 把特有变量具体的值替换进去
                        cond = cond.replace(lv.toString(), assignList.get(lv.toString()).trim());
                    }//不包含特有变量的条件呢？？？
//                    else stepConditions.add(cond);
                }

                // 去掉各种方法调用
                // virtualinvoke：用于调用普通的成员方法，进行virtual dispatch
                if (cond.contains("virtualinvoke ")) {
                    cond = cond.replace("virtualinvoke ", "");
                    int lAngleBrac = cond.indexOf("<");
                    int rAngleBrac = matchBracket(cond, lAngleBrac);
                    // 调用的方法名称，全名
                    String invokeMeth = cond.substring(lAngleBrac+1, rAngleBrac);
                    // 处理这个调用方法的参数
                    int lPara = rAngleBrac + 1;
                    int rPara = matchBracket(cond, lPara);
                    String invoPara = cond.substring(lPara+1, rPara);
                    // 去掉调用后的语句
                    cond = cond.replace(".<"+invokeMeth+">", "");
                    cond = cond.replace("("+invoPara+")", "");
                    InvokeStmt invokeStmt = new InvokeStmt("virtualinvoke", invokeMeth, invoPara, cond);
                    invokeStmts.add(invokeStmt);
                }
                stepConditions.add(cond);
            }
        } else {
            stepConditions = stepConditionsWithJimpleVars;
        }
        return stepConditions;
    }

    // 匹配括号，无论是圆括号还是尖括号
    private int matchBracket(String waitDetectStr, int lBracketsIndex) {
        int rBracketsIndex = -1;

        // 检测传进来的左括号索引是否超出
        if (waitDetectStr.length() > lBracketsIndex) {
            // 检测传进来的索引对应的字符是否为左括号
            if ('(' == waitDetectStr.charAt(lBracketsIndex)) {
                // 获取另一配对括号位置
                int count = 1;
                for (int i = lBracketsIndex + 1; i < waitDetectStr.length(); i++) {
                    if (waitDetectStr.charAt(i) == '(') {
                        count++;
                    } else if (waitDetectStr.charAt(i) == ')') {
                        count--;
                    }
                    if (count == 0) {
                        rBracketsIndex = i;
                        break;
                    }
                }
            }else if ('<' == waitDetectStr.charAt(lBracketsIndex)) {
                // 获取另一配对括号位置
                int count = 1;
                for (int i = lBracketsIndex + 1; i < waitDetectStr.length(); i++) {
                    if (waitDetectStr.charAt(i) == '<') {
                        count++;
                    } else if (waitDetectStr.charAt(i) == '>') {
                        count--;
                    }
                    if (count == 0) {
                        rBracketsIndex = i;
                        break;
                    }
                }
            }
        }
        return rBracketsIndex;
    }

    public String solve(String pathConstraint) throws Exception {
        // 直接用z3来解析
        return Z3Solver.solve(pathConstraint, converJimpleVars);
    }

    // 随机生成一条测试语句 **=**
    public String randomTC(List<Local> parameters) {

        String varName;
        String varValue = "";
        String testinput = "";

        for (Local para : parameters) {
            varName = para.getName();
            if ("int".equals(para.getType().toString())) {
                varValue = String.valueOf((int)(Math.random() * 10));
            }
            if ("String".equals(para.getType().toString())) {
                varValue = "abc";
            }
            //其它的基本类型没写
            // 要实现的！！！！！！！！！
            testinput = testinput + " " + varName + "=" + varValue;
        }
        return testinput;
    }

}
